import tensorflow as tf
from tensorflow import keras
from keras.applications import ResNet50
from keras.layers import Dense, GlobalAveragePooling2D
from keras.models import Model
import numpy as np

import cv2,os

# 1. 加载数据
train_dir = r'ai/train'
val_dir = r'ai/val'
test_dir=r'ai/test'
img_size = (224, 224)  # ResNet50 输入尺寸
batch_size = 12


def convert_pic(image_path:str,h1=224,w1=224):
    img=cv2.imread(image_path)
    p1=r'd:\1.jpg'
    index=image_path.rindex('.')
    if index>0:
        p1=f'{image_path[0:index]}_copy{image_path[index:]}'    

    h,w=img.shape[:2]

    # img.resize(h1,w1)
    if h!=h1 or w!=w:
        img1=cv2.resize(img,(h1,w1))
    else:
        img1=img

    cv2.imwrite(p1,img1)

    pass


def convert_pic_dir(dir:str):
    os.chdir(dir)
    for i in os.listdir(dir):
        convert_pic(f'{dir}/{i}')



# convert_pic_dir(f'{os.path.dirname(__file__)}/{test_dir}')
# exit()

train_ds = keras.utils.image_dataset_from_directory(
    train_dir,
    image_size=img_size,
    batch_size=batch_size,
    shuffle=True
)

val_ds = keras.utils.image_dataset_from_directory(
    val_dir,
    image_size=img_size,
    batch_size=3,
    shuffle=False
)

# 2. 数据预处理（归一化到 [-1, 1]，匹配 ResNet 预训练参数）
preprocess_input = keras.applications.resnet50.preprocess_input

# 定义归一化函数
def normalize_img(image, label):
    return tf.cast(image, tf.float32) / 255.0, label

# 应用归一化
ds_dataset = train_ds.map(normalize_img)
val_dataset = val_ds.map(normalize_img)

# 3. 加载预训练模型（冻结基础层）
base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
base_model.trainable = False  # 冻结预训练权重

# 4. 添加分类头
x = base_model.output
x = GlobalAveragePooling2D()(x)  # 全局平均池化
x = Dense(128, activation='relu')(x)  # 全连接层
predictions = Dense(len(train_ds.class_names), activation='softmax')(x)  # 输出层（类别数）
# predictions = Dense(2, activation='softmax')(x)  # 输出层（类别数）

model = Model(inputs=base_model.input, outputs=predictions)

# 5. 编译模型
model.compile(
    optimizer=keras.optimizers.Adam(learning_rate=0.001),
    loss='sparse_categorical_crossentropy',
    metrics=['accuracy']
)

# 6. 训练模型
history = model.fit(
    train_ds,
    validation_data=val_ds,
    epochs=10  # 可根据效果调整
)

# 7. 预测新图片
def predict_image(img_path):
    img = keras.utils.load_img(img_path, target_size=img_size)
    img_array = keras.utils.img_to_array(img)
    img_array = preprocess_input(img_array[np.newaxis, ...])  # 增加批次维度
    pred = model.predict(img_array)
    class_idx = pred.argmax()
    return train_ds.class_names[class_idx]

# 测试预测
# print(predict_image(r'd:\car.jpg'))  # 输出预测类别

for i in os.listdir(test_dir):
    tp=f'{os.path.dirname(__file__)}/{test_dir}/{i}'
    print(i,predict_image(tp))
